/*
 * Copyright (c) 2010 Google Inc.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

package com.google.api.client.apache;

import org.apache.http.annotation.ThreadSafe;
import org.apache.http.conn.ConnectTimeoutException;
import org.apache.http.conn.scheme.LayeredSchemeSocketFactory;
import org.apache.http.conn.ssl.AllowAllHostnameVerifier;
import org.apache.http.conn.ssl.BrowserCompatHostnameVerifier;
import org.apache.http.conn.ssl.StrictHostnameVerifier;
import org.apache.http.conn.ssl.TrustStrategy;
import org.apache.http.conn.ssl.X509HostnameVerifier;
import org.apache.http.params.HttpConnectionParams;
import org.apache.http.params.HttpParams;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.net.Socket;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;
import java.security.KeyManagementException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.UnrecoverableKeyException;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

@ThreadSafe
public class SSLSocksSocketFactory implements LayeredSchemeSocketFactory {

  public static final String TLS = "TLS";
  public static final String SSL = "SSL";
  public static final String SSLV2 = "SSLv2";

  public static final X509HostnameVerifier ALLOW_ALL_HOSTNAME_VERIFIER =
      new AllowAllHostnameVerifier();

  public static final X509HostnameVerifier BROWSER_COMPATIBLE_HOSTNAME_VERIFIER =
      new BrowserCompatHostnameVerifier();

  public static final X509HostnameVerifier STRICT_HOSTNAME_VERIFIER = new StrictHostnameVerifier();

  /**
   * Gets the default factory, which uses the default JVM settings for secure connections.
   * 
   * @return the default factory
   */
  public static SSLSocksSocketFactory getSocketFactory() {
    return new SSLSocksSocketFactory();
  }

  private final javax.net.ssl.SSLSocketFactory socketfactory;
  private final X509HostnameVerifier hostnameVerifier;

  private static SSLContext createSSLContext(String algorithm, final KeyStore keystore,
      final String keystorePassword, final KeyStore truststore, final SecureRandom random,
      final TrustStrategy trustStrategy) throws NoSuchAlgorithmException, KeyStoreException,
      UnrecoverableKeyException, KeyManagementException {
    if (algorithm == null) {
      algorithm = TLS;
    }
    KeyManagerFactory kmfactory =
        KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
    kmfactory.init(keystore, keystorePassword != null ? keystorePassword.toCharArray() : null);
    KeyManager[] keymanagers = kmfactory.getKeyManagers();
    TrustManagerFactory tmfactory =
        TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
    tmfactory.init(truststore);
    TrustManager[] trustmanagers = tmfactory.getTrustManagers();
    if (trustmanagers != null && trustStrategy != null) {
      for (int i = 0; i < trustmanagers.length; i++) {
        TrustManager tm = trustmanagers[i];
        if (tm instanceof X509TrustManager) {
          trustmanagers[i] = new TrustManagerDecorator((X509TrustManager) tm, trustStrategy);
        }
      }
    }

    SSLContext sslcontext = SSLContext.getInstance(algorithm);
    sslcontext.init(keymanagers, trustmanagers, random);
    return sslcontext;
  }

  private static SSLContext createDefaultSSLContext() {
    try {
      return createSSLContext(TLS, null, null, null, null, null);
    } catch (Exception ex) {
      throw new IllegalStateException("Failure initializing default SSL context", ex);
    }
  }

  public SSLSocksSocketFactory(final SSLContext sslContext) {
    this(sslContext, BROWSER_COMPATIBLE_HOSTNAME_VERIFIER);
  }

  public SSLSocksSocketFactory(final SSLContext sslContext,
      final X509HostnameVerifier hostnameVerifier) {
    super();
    this.socketfactory = sslContext.getSocketFactory();
    this.hostnameVerifier = hostnameVerifier;
  }

  private SSLSocksSocketFactory() {
    this(createDefaultSSLContext());
  }

  /**
   * @param params If the parameters contain values for socks.host and socks.port they are used as
   *        the socket server. If the values are not set (or set to <code>null</code>) the function
   *        simply returns a new instance of {@link Socket} class using default constructor.
   * 
   * @return the new socket
   */
  public Socket createSocket(final HttpParams params) {
    Socket sock;
    if (params == null) {
      sock = new Socket();
    } else {
      String proxyHost = (String) params.getParameter("socks.host");
      Integer proxyPort = (Integer) params.getParameter("socks.port");
      if (proxyHost != null && proxyPort != null) {
        InetSocketAddress socksaddr = new InetSocketAddress(proxyHost, proxyPort);
        Proxy proxy = new Proxy(Proxy.Type.SOCKS, socksaddr);
        sock = new Socket(proxy);
      } else {
        sock = new Socket();
      }
    }
    return sock;
  }

  public Socket connectSocket(final Socket socket, final InetSocketAddress remoteAddress,
      final InetSocketAddress localAddress, final HttpParams params) throws IOException,
      UnknownHostException, ConnectTimeoutException {
    if (remoteAddress == null) {
      throw new IllegalArgumentException("Remote address may not be null");
    }
    if (params == null) {
      throw new IllegalArgumentException("HTTP parameters may not be null");
    }
    Socket sock = socket != null ? socket : createSocket(params);
    if (localAddress != null) {
      sock.setReuseAddress(HttpConnectionParams.getSoReuseaddr(params));
      sock.bind(localAddress);
    }

    int connTimeout = HttpConnectionParams.getConnectionTimeout(params);
    int soTimeout = HttpConnectionParams.getSoTimeout(params);

    try {
      sock.connect(remoteAddress, connTimeout);
    } catch (SocketTimeoutException ex) {
      throw new ConnectTimeoutException("Connect to " + remoteAddress.getHostName() + "/"
          + remoteAddress.getAddress() + " timed out");
    }
    sock.setSoTimeout(soTimeout);
    SSLSocket sslsock;
    // Setup SSL layering if necessary
    if (sock instanceof SSLSocket) {
      sslsock = (SSLSocket) sock;
    } else {
      sslsock =
          (SSLSocket) this.socketfactory.createSocket(sock, remoteAddress.getHostName(),
              remoteAddress.getPort(), true);
    }
    if (this.hostnameVerifier != null) {
      try {
        this.hostnameVerifier.verify(remoteAddress.getHostName(), sslsock);
        // verifyHostName() didn't blowup - good!
      } catch (IOException iox) {
        // close the socket before re-throwing the exception
        try {
          sslsock.close();
        } catch (Exception x) { /* ignore */
        }
        throw iox;
      }
    }
    return sslsock;
  }


  /**
   * Checks whether a socket connection is secure. This factory creates TLS/SSL socket connections
   * which, by default, are considered secure. <br/>
   * Derived classes may override this method to perform runtime checks, for example based on the
   * cypher suite.
   * 
   * @param sock the connected socket
   * 
   * @return <code>true</code>
   * 
   * @throws IllegalArgumentException if the argument is invalid
   */
  public boolean isSecure(final Socket sock) throws IllegalArgumentException {
    if (sock == null) {
      throw new IllegalArgumentException("Socket may not be null");
    }
    // This instanceof check is in line with createSocket() above.
    if (!(sock instanceof SSLSocket)) {
      throw new IllegalArgumentException("Socket not created by this factory");
    }
    // This check is performed last since it calls the argument object.
    if (sock.isClosed()) {
      throw new IllegalArgumentException("Socket is closed");
    }
    return true;
  }

  public Socket createLayeredSocket(final Socket socket, final String host, final int port,
      final boolean autoClose) throws IOException, UnknownHostException {
    SSLSocket sslSocket =
        (SSLSocket) this.socketfactory.createSocket(socket, host, port, autoClose);
    if (this.hostnameVerifier != null) {
      this.hostnameVerifier.verify(host, sslSocket);
    }
    // verifyHostName() didn't blowup - good!
    return sslSocket;
  }

  public X509HostnameVerifier getHostnameVerifier() {
    return this.hostnameVerifier;
  }
}
